import tensorflow as tf
import numpy as np
from layer1_supervised_lenet import inference
from foolbox.models import TensorFlowModel
from foolbox.criteria import Misclassification
from foolbox.attacks import DeepFoolAttack
from tensorflow.contrib.learn.python.learn.datasets.mnist import read_data_sets

mnist = read_data_sets("MNIST_data/",one_hot=True)

images = tf.placeholder(tf.float32,shape=(None,784))
y = tf.placeholder(tf.float32, [None, 10])
loss,logits = inference(images,y,keep_prob=1.0)
saver = tf.train.Saver()

with tf.Session() as session:
    saver.restore(session, tf.train.latest_checkpoint('./finetune_model/model/layer1'))
    print('finish loading model!')

    model = TensorFlowModel(images, logits, bounds=(0, 1))
    criterion = Misclassification()
    attack = DeepFoolAttack(model,criterion)
    per_dataset = np.zeros((1000,784))
    for i in range(1000):
        image = mnist.test.images[i]
        label = np.argmax(model.predictions(image))
        T = np.argmax(mnist.test.labels[i])
        adversarial = attack(image,label)
        per_dataset[i] = adversarial
        print("making" + str(i) + "picture")
        print(label,T)
    np.save("deepfoolattack_layer1_lenet-5",per_dataset)


